///
/// Copyright (C) 2016, Dependable Systems Laboratory, EPFL
/// Copyright (C) 2016, Cyberhaven
///
/// Permission is hereby granted, free of charge, to any person obtaining a copy
/// of this software and associated documentation files (the "Software"), to deal
/// in the Software without restriction, including without limitation the rights
/// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
/// copies of the Software, and to permit persons to whom the Software is
/// furnished to do so, subject to the following conditions:
///
/// The above copyright notice and this permission notice shall be included in all
/// copies or substantial portions of the Software.
///
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
/// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
/// SOFTWARE.
///

#ifndef S2E_PLUGINS_RecipeDescriptor_H_
#define S2E_PLUGINS_RecipeDescriptor_H_

#include <s2e/Plugins/OSMonitors/OSMonitor.h>
#include <s2e/Plugins/VulnerabilityAnalysis/PovGenerator.h>
#include <s2e/cpu.h>

#include <klee/util/ExprUtil.h>
#include <klee/util/Ref.h>
#include <llvm/Support/Path.h>
#include <vector>

#include "Register.h"

namespace s2e {
namespace plugins {
namespace recipe {

typedef pov::PovOptions PovOptions;
typedef pov::PovType PovType;
typedef pov::VariableRemapping VariableRemapping;
typedef pov::ExprList ExprList;

typedef std::vector<klee::ref<Register>> RegList;

// Left part of expression
class Left {
public:
    enum Type {
        INV,         // invalid
        REGBYTE,     // register[offset] (EAX[0])
        REGPTR,      // memory referenced by register+offset
        REGPTR_EXEC, // register must point to executable memory
        REGPTR_PTR,  // [register+offs1][offs2] ([ESP+4][0])
        ADDR,        // memory referenced by address
    };

private:
    Type m_type = INV;

    klee::ref<Register> m_reg;
    off_t m_offset = 0;

    uint64_t m_addr = 0;
    off_t m_mem_offset = 0;

    Left() {
    }

public:
    // These are required for klee::ref<>
    unsigned refCount = 0;
    bool permanent = false;

    // Access part of a register: register[offset] (EAX[0])
    static klee::ref<Left> createRegByte(klee::ref<Register> reg) {
        Left *ret = new Left();
        ret->m_type = REGBYTE;
        ret->m_reg = reg;
        return klee::ref<Left>(ret);
    }

    // Access memory: [reg + offset]
    static klee::ref<Left> createRegPtr(klee::ref<Register> reg, off_t offset) {
        if (reg->idx() > 0) {
            return nullptr;
        }

        Left *ret = new Left();
        ret->m_type = REGPTR;
        ret->m_reg = reg;
        ret->m_offset = offset;
        return klee::ref<Left>(ret);
    }

    // Register must point to executable memory
    static klee::ref<Left> createRegPtrExec(klee::ref<Register> reg) {
        if (reg->idx() > 0) {
            return nullptr;
        }

        Left *ret = new Left();
        ret->m_type = REGPTR_EXEC;
        ret->m_reg = reg;
        return klee::ref<Left>(ret);
    }

    // [register+offs1][offs2] ([ESP+4][0])
    static klee::ref<Left> createRegPtrPtr(klee::ref<Register> reg, off_t offset, off_t mem_offset) {
        Left *ret = new Left();
        ret->m_type = REGPTR_PTR;
        ret->m_reg = reg;
        ret->m_offset = offset;
        ret->m_mem_offset = mem_offset;
        return klee::ref<Left>(ret);
    }

    // [address]
    static klee::ref<Left> createAddr(uint64_t addr) {
        Left *ret = new Left();
        ret->m_type = ADDR;
        ret->m_addr = addr;
        return klee::ref<Left>(ret);
    }

    template <typename T> friend T &operator<<(T &stream, const Left &l) {
        switch (l.m_type) {
            case REGBYTE:
                stream << l.m_reg->name() << "[" << hexval(l.m_reg->idx()) << "]";
                break;
            case ADDR:
                stream << "[" << hexval(l.m_addr) << "]";
                break;
            case REGPTR:
                stream << "[" << l.m_reg->name() << "+" << hexval(l.m_offset) << "]";
                break;
            case REGPTR_EXEC:
                stream << "*" << l.m_reg->name() << " points to executable memory*";
                break;
            case REGPTR_PTR:
                stream << "[" << l.m_reg->name() << "+" << hexval(l.m_offset) << "][" << hexval(l.m_mem_offset) << "]";
                break;
            default:
                stream << "INV";
                break;
        }
        return stream;
    }

    std::string name() const {
        std::ostringstream ss;
        ss << *this;
        return ss.str();
    }

    Type type() const {
        return m_type;
    }
    const klee::ref<Register> &reg() {
        return m_reg;
    }
    uint64_t addr() const {
        return m_addr;
    }
    off_t offset() const {
        return m_offset;
    }
    off_t memOffset() const {
        return m_mem_offset;
    }
};

// Right part of expression.
class Right {
public:
    enum Type {
        INV,        // invalid
        REGBYTE,    // register[offset] (EAX[0])
        NEGOTIABLE, // can take arbitrary values
        CONCRETE    // has concrete value
    };

private:
    Type m_type = INV;

    std::string m_tag = "INVALID";
    klee::ref<Register> m_reg;

    uint64_t m_value = 0;
    klee::Expr::Width m_valueWidth = klee::Expr::InvalidWidth;

    Right() {
    }

public:
    // These are required for klee::ref<>
    unsigned refCount = 0;
    bool permanent = false;

    static klee::ref<Right> createRegByte(const klee::ref<Register> &reg) {
        Right *ret = new Right();
        ret->m_type = REGBYTE;
        ret->m_reg = reg;
        return klee::ref<Right>(ret);
    }

    static klee::ref<Right> createNegotiable(const std::string &tag) {
        Right *ret = new Right();
        ret->m_type = NEGOTIABLE;
        ret->m_tag = tag;
        return klee::ref<Right>(ret);
    }

    static klee::ref<Right> createConcrete(uint64_t value, klee::Expr::Width width) {
        Right *ret = new Right();
        ret->m_type = CONCRETE;
        ret->m_valueWidth = width;
        ret->m_value = value;

        s2e_assert(nullptr, value == klee::bits64::truncateToNBits(value, width),
                   "Value " << hexval(value) << " does not fit " << width << " bits");

        return klee::ref<Right>(ret);
    }

    static klee::ref<Right> createInvalid() {
        return klee::ref<Right>(new Right());
    }

    template <typename T> friend T &operator<<(T &stream, const Right &r) {
        switch (r.m_type) {
            case INV:
                stream << "INV";
                break;
            case REGBYTE:
                stream << r.m_reg->name() << "[" << hexval(r.m_reg->idx()) << "]";
                break;
            case NEGOTIABLE:
                stream << r.m_tag;
                break;
            case CONCRETE:
                stream << hexval(r.m_value);
                break;
        }
        return stream;
    }

    std::string name() const {
        std::ostringstream ss;
        ss << *this;
        return ss.str();
    }

    Type type() const {
        return m_type;
    }
    klee::Expr::Width valueWidth() const {
        return m_valueWidth;
    }
    uint64_t value() const {
        return m_value;
    }
    const klee::ref<Register> &reg() {
        return m_reg;
    }
};

// Precondtion is expressed as "left == right"
struct Precondition {
    klee::ref<Left> left;
    klee::ref<Right> right;

    Precondition() {
    }
    Precondition(const klee::ref<Left> &l, const klee::ref<Right> &r) : left(l), right(r) {
    }

    template <typename T> friend T &operator<<(T &stream, const Precondition &p) {
        stream << p.left;
        if (p.left->type() != Left::REGPTR_EXEC) {
            stream << " == " << p.right;
        }
        return stream;
    }
};

typedef std::vector<Precondition> Preconditions;

enum RecipeArch { RECIPE_I386, RECIPE_AMD64 };

enum RecipePlatform { RECIPE_GENERIC, RECIPE_DECREE };

struct RecipeSettings {
    PovType type;
    klee::ref<Register> gp;
    uint32_t regMask;
    uint32_t ipMask;
    uint32_t skip;
    std::string moduleName;
    RecipePlatform platform;
    RecipeArch arch;

    RecipeSettings() : type(PovType::POV_GENERAL), regMask(), ipMask(), skip(), moduleName() {
    }
    RecipeSettings(PovType type, klee::ref<Register> gp, uint32_t regMask, uint32_t ipMask, uint32_t skip,
                   std::string moduleName)
        : type(type), gp(gp), regMask(regMask), ipMask(ipMask), skip(skip), moduleName(moduleName) {
    }

    RecipeSettings(PovType type, const std::string gp)
        : RecipeSettings(type, Register::fromName(gp, 0), 0xffffffff, 0xffffffff, 0, "") {
    }

    RecipeSettings(PovType type, const std::string gp, uint32_t regMask, uint32_t ipMask)
        : RecipeSettings(type, Register::fromName(gp, 0), regMask, ipMask, 0, "") {
    }
};

class RecipeConditions {
public:
    ExprList constraints;
    ExprList usedExprs;
    VariableRemapping remappings;
    RegList usedRegs;

    RecipeConditions &operator=(const RecipeConditions &other) {
        if (this == &other) {
            return *this;
        }

        this->constraints = other.constraints;
        this->usedExprs = other.usedExprs;
        this->remappings = other.remappings;
        this->usedRegs = other.usedRegs;

        return *this;
    }
};

enum EIPType {
    SYMBOLIC_EIP = 1u << 0,
    CONCRETE_EIP = 1u << 1,
};

struct StateConditions {
    ModuleDescriptor module;
    klee::ref<klee::Expr> nextEip;
    EIPType eipType;
};

struct RecipeDescriptor {
    RecipeSettings settings;
    Preconditions preconditions;
    EIPType eipType;

    static RecipeDescriptor *fromFile(const std::string &recipeFile);
    static bool mustTryRecipe(const RecipeDescriptor &recipe, const std::string &recipeName, const StateConditions &sc,
                              uint64_t eip);

    bool isUsable(S2EExecutionState *state, OSMonitor *monitor) const;

private:
    uint64_t concreteTargetEIP;
    bool parseSettingsLine(const std::string &line);
    bool parsePreconditionLine(const std::string &line);
    bool isValid() const;
};
} // namespace recipe
} // namespace plugins
} // namespace s2e

#endif
